import bisect
import gc
import glob
import random

import torch
from transformers import BertTokenizer
from others.logging import logger

import numpy as np


class Batch(object):
    def _pad(self, data, pad_id, width=-1, multi=-1):
        # print(data)
        if (width == -1):
            if multi == -1:
                width = max(len(d) for d in data)
            else:
                multi = max(len(d) for d in data)
                width = max(max(len(l) for l in d) if len(d) > 0 else 0 for d in data)
        # print(width)
        if multi == -1:
            rtn_data = [d + [pad_id] * (width - len(d)) for d in data]
        else:
            # print(width)
            rtn_data = []
            # print('data:')
            # print(data)
            for d in data:
                # print('d')
                # print(d)
                rtn_l = []
                for l in d:
                    # print('l')
                    # print(l)
                    if width > len(l):
                        rtn_l.append(l + [pad_id] * (width - len(l)))
                    else:
                        rtn_l.append(l[:width])
                rtn_data.append(rtn_l + [[0] * width] * (multi - len(rtn_l)))
            # rtn_data = [[l + [pad_id] * (width - len(l)) if len(l)>0 else [pad_id]*width for l in d]for d in data]
            # print(rtn_data)
            # print(len(rtn_data))
            # for d in rtn_data:
            #     print('d:' + str(len(d)))
            # print(np.array(rtn_data).shape)
        return rtn_data

    def __init__(self, data=None, device=None, is_test=False):
        """Create a Batch from a list of examples."""
        if data is not None:
            self.batch_size = len(data)
            # pre_src = [x[0] for x in data]  # fact, law[], evidence
            # print('src')
            pre_fact = [x[0][0] for x in data]
            pre_law = [x[0][1] for x in data]
            pre_evidence = [x[0][2] for x in data]

            # print('segs')
            pre_tgt = [x[1] for x in data]  # claim
            # pre_segs = [x[2] for x in data]  # fsegs, lsegs, esegs
            pre_fsegs = [x[2][0] for x in data]
            pre_lsegs = [x[2][1] for x in data]
            pre_esegs = [x[2][2] for x in data]

            # print('fclss')
            # pre_clss = [x[3] for x in data]  # fclss , (lclss,eclss)
            pre_fclss = [x[3][0] for x in data]
            pre_lclss = [x[3][1] for x in data]
            pre_eclss = [x[3][2] for x in data]

            pre_src_sent_labels = [x[4] for x in data]  # unuse

            # fact, law, evidence = pre_src
            fact = torch.tensor(self._pad(pre_fact, 0))
            law = torch.tensor(self._pad(pre_law, 0, multi=1))
            evidence = torch.tensor(self._pad(pre_evidence, 0))

            tgt = torch.tensor(self._pad(pre_tgt, 0))

            fsegs = torch.tensor(self._pad(pre_fsegs, 0))
            lsegs = torch.tensor(self._pad(pre_lsegs, 0, multi=1))
            esegs = torch.tensor(self._pad(pre_esegs, 0))

            mask_src = 1 - (fact == 0)
            mask_tgt = 1 - (tgt == 0)

            fclss = torch.tensor(self._pad(pre_fclss, -1))
            lclss = torch.tensor(self._pad(pre_lclss, -1))
            eclss = torch.tensor(self._pad(pre_eclss, -1))

            src_sent_labels = torch.tensor(self._pad(pre_src_sent_labels, 0))

            mask_cls = 1 - (fclss == -1)
            fclss[fclss == -1] = 0

            setattr(self, 'clss', fclss.to(device))
            setattr(self, 'mask_cls', mask_cls.to(device))
            setattr(self, 'src_sent_labels', src_sent_labels.to(device))

            src = [fact, law, evidence]
            segs = [fsegs, lsegs, esegs]

            setattr(self, 'fact', fact.to(device))
            setattr(self, 'law', law.to(device))
            setattr(self, 'evidence', evidence.to(device))

            setattr(self, 'fsegs', fsegs.to(device))
            setattr(self, 'lsegs', lsegs.to(device))
            setattr(self, 'esegs', esegs.to(device))

            setattr(self, 'tgt', tgt.to(device))
            # setattr(self, 'segs', segs.to(device))
            setattr(self, 'mask_src', mask_src.to(device))
            setattr(self, 'mask_tgt', mask_tgt.to(device))

            if (is_test):
                src_str = [[x[-4]] + [x[-3]] + [x[-2]] for x in data]
                setattr(self, 'src_str', src_str)
                tgt_str = [x[-1] for x in data]
                setattr(self, 'tgt_str', tgt_str)
            # src = torch.tensor(self._pad(pre_src, 0))
            # tgt = torch.tensor(self._pad(pre_tgt, 0))
            #
            # segs = torch.tensor(self._pad(pre_segs, 0))
            # mask_src = 1 - (src == 0)
            # mask_tgt = 1 - (tgt == 0)
            #
            # clss = torch.tensor(self._pad(pre_clss, -1))
            # src_sent_labels = torch.tensor(self._pad(pre_src_sent_labels, 0))
            # mask_cls = 1 - (clss == -1)
            # clss[clss == -1] = 0
            # setattr(self, 'clss', clss.to(device))
            # setattr(self, 'mask_cls', mask_cls.to(device))
            # setattr(self, 'src_sent_labels', src_sent_labels.to(device))
            #
            # setattr(self, 'src', src.to(device))
            # setattr(self, 'tgt', tgt.to(device))
            # setattr(self, 'segs', segs.to(device))
            # setattr(self, 'mask_src', mask_src.to(device))
            # setattr(self, 'mask_tgt', mask_tgt.to(device))

            # if (is_test):
            #     src_str = [x[-2] for x in data]
            #     setattr(self, 'src_str', src_str)
            #     tgt_str = [x[-1] for x in data]
            #     setattr(self, 'tgt_str', tgt_str)

    def __len__(self):
        return self.batch_size


def load_dataset(args, corpus_type, shuffle):
    """
    Dataset generator. Don't do extra stuff here, like printing,
    because they will be postponed to the first loading time.

    Args:
        corpus_type: 'train' or 'valid'
    Returns:
        A list of dataset, the dataset(s) are lazily loaded.
    """
    assert corpus_type in ["train", "valid", "test"]

    def _lazy_dataset_loader(pt_file, corpus_type):
        dataset = torch.load(pt_file)
        logger.info('Loading %s dataset from %s, number of examples: %d' %
                    (corpus_type, pt_file, len(dataset)))
        return dataset

    # Sort the glob output by file name (by increasing indexes).
    pts = sorted(glob.glob(args.bert_data_path + '.' + corpus_type + '.[0-9]*.pt'))
    if pts:
        if (shuffle):
            random.shuffle(pts)

        for pt in pts:
            yield _lazy_dataset_loader(pt, corpus_type)
    else:
        # Only one inputters.*Dataset, simple!
        pt = args.bert_data_path + '.' + corpus_type + '.pt'
        yield _lazy_dataset_loader(pt, corpus_type)


def abs_batch_size_fn(new, count):
    src, tgt = new[0], new[1]
    global max_n_sents, max_n_tokens, max_size
    if count == 1:
        max_size = 0
        max_n_sents = 0
        max_n_tokens = 0
    max_n_sents = max(max_n_sents, len(tgt))
    max_size = max(max_size, max_n_sents)
    src_elements = count * max_size
    if (count > 6):
        return src_elements + 1e3
    return src_elements


def ext_batch_size_fn(new, count):
    if (len(new) == 4):
        pass
    src, labels = new[0], new[4]
    global max_n_sents, max_n_tokens, max_size
    if count == 1:
        max_size = 0
        max_n_sents = 0
        max_n_tokens = 0
    max_n_sents = max(max_n_sents, len(src))
    max_size = max(max_size, max_n_sents)
    src_elements = count * max_size
    return src_elements


class Dataloader(object):
    def __init__(self, args, datasets, batch_size,
                 device, shuffle, is_test):
        self.args = args
        self.datasets = datasets
        self.batch_size = batch_size
        self.device = device
        self.shuffle = shuffle
        self.is_test = is_test
        self.cur_iter = self._next_dataset_iterator(datasets)
        assert self.cur_iter is not None

    def __iter__(self):
        dataset_iter = (d for d in self.datasets)
        while self.cur_iter is not None:
            for batch in self.cur_iter:
                yield batch
            self.cur_iter = self._next_dataset_iterator(dataset_iter)

    def _next_dataset_iterator(self, dataset_iter):
        try:
            # Drop the current dataset for decreasing memory
            if hasattr(self, "cur_dataset"):
                self.cur_dataset = None
                gc.collect()
                del self.cur_dataset
                gc.collect()

            self.cur_dataset = next(dataset_iter)
        except StopIteration:
            return None

        return DataIterator(args=self.args,
                            dataset=self.cur_dataset, batch_size=self.batch_size,
                            device=self.device, shuffle=self.shuffle, is_test=self.is_test)


class DataIterator(object):
    def __init__(self, args, dataset, batch_size, device=None, is_test=False,
                 shuffle=True):
        self.args = args
        self.batch_size, self.is_test, self.dataset = batch_size, is_test, dataset
        self.iterations = 0
        self.device = device
        self.shuffle = shuffle

        self.sort_key = lambda x: len(x[1])

        self._iterations_this_epoch = 0
        if (self.args.task == 'abs'):
            self.batch_size_fn = abs_batch_size_fn
        else:
            self.batch_size_fn = ext_batch_size_fn

    def data(self):
        if self.shuffle:
            random.shuffle(self.dataset)
        xs = self.dataset
        return xs

    #################################################
    def preprocess(self, ex, is_test):
        fact, law, evidence, claim = ex['fact'], ex['law'], ex['evidence'], ex['claim']
        fsegs, lsegs, esegs = ex['fsegs'], ex['lsegs'], ex['esegs']
        fclss, lclss, eclss = ex['fclss'], ex['lclss'], ex['eclss']
        ftxt, ltxt, etxt, ctxt = ex['ftxt'], ex['ltxt'], ex['etxt'], ex['ctxt']

        tgt = claim[:self.args.max_tgt_len][:-1] + [2]

        # print(tgt)
        # tokenizer = BertTokenizer.from_pretrained('hfl/chinese-bert-wwm', do_lower_case=True)
        # tgtt = tokenizer.convert_ids_to_tokens(list(tgt))
        # print(tgtt)
        # print(ctxt)

        if not self.args.use_interval:
            fsegs = [0] * len(fsegs)
            lsegs = [[0] * len(lsegs) for i in range(len(law))]
            esegs = [0] * len(esegs)

        end_id = [fact[-1]]
        fact = fact[:-1][:self.args.max_pos - 1] + end_id
        law = [l[:-1][:self.args.max_pos - 1] + end_id for l in law]
        evidence = evidence[:-1][:self.args.max_pos - 1] + end_id

        fsegs = fsegs[:self.args.max_pos]
        lsegs = [l[:self.args.max_pos] for l in lsegs]
        esegs = esegs[:self.args.max_pos]

        # src = ex['src']
        # tgt = ex['tgt'][:self.args.max_tgt_len][:-1] + [2]
        # src_sent_labels = ex['src_sent_labels']
        # segs = ex['segs']
        # if (not self.args.use_interval):
        #     segs = [0] * len(segs)
        # clss = ex['clss']
        # src_txt = ex['src_txt']
        # tgt_txt = ex['tgt_txt']
        #
        # end_id = [src[-1]]
        # src = src[:-1][:self.args.max_pos - 1] + end_id
        # segs = segs[:self.args.max_pos]
        # max_sent_id = bisect.bisect_left(clss, self.args.max_pos)
        # src_sent_labels = src_sent_labels[:max_sent_id]
        # clss = clss[:max_sent_id]
        # src_txt = src_txt[:max_sent_id]

        src = [fact, law, evidence]
        segs = [fsegs, lsegs, esegs]
        clss = [fclss, lclss, eclss]
        src_txt = [ftxt, ltxt, etxt]
        tgt_txt = ctxt

        src_sent_labels = [0] * len(fact)
        if (is_test):
            return src, tgt, segs, clss, src_sent_labels, src_txt, tgt_txt
        else:
            return src, tgt, segs, clss, src_sent_labels

    def batch_buffer(self, data, batch_size):
        minibatch, size_so_far = [], 0
        for ex in data:
            if (len(ex['fact']) == 0):
                continue
            ex = self.preprocess(ex, self.is_test)
            if (ex is None):
                continue
            minibatch.append(ex)
            size_so_far = self.batch_size_fn(ex, len(minibatch))
            if size_so_far == batch_size:
                yield minibatch
                minibatch, size_so_far = [], 0
            elif size_so_far > batch_size:
                yield minibatch[:-1]
                minibatch, size_so_far = minibatch[-1:], self.batch_size_fn(ex, 1)
        if minibatch:
            yield minibatch

    def batch(self, data, batch_size):
        """Yield elements from data in chunks of batch_size."""
        minibatch, size_so_far = [], 0
        for ex in data:
            minibatch.append(ex)
            size_so_far = self.batch_size_fn(ex, len(minibatch))
            if size_so_far == batch_size:
                yield minibatch
                minibatch, size_so_far = [], 0
            elif size_so_far > batch_size:
                yield minibatch[:-1]
                minibatch, size_so_far = minibatch[-1:], self.batch_size_fn(ex, 1)
        if minibatch:
            yield minibatch

    def create_batches(self):
        """ Create batches """
        data = self.data()
        for buffer in self.batch_buffer(data, self.batch_size * 300):

            if (self.args.task == 'abs'):
                p_batch = sorted(buffer, key=lambda x: len(x[2]))
                p_batch = sorted(p_batch, key=lambda x: len(x[1]))
            else:
                p_batch = sorted(buffer, key=lambda x: len(x[2]))

            p_batch = self.batch(p_batch, self.batch_size)

            p_batch = list(p_batch)
            if (self.shuffle):
                random.shuffle(p_batch)
            for b in p_batch:
                if (len(b) == 0):
                    continue
                yield b

    def __iter__(self):
        while True:
            self.batches = self.create_batches()
            for idx, minibatch in enumerate(self.batches):
                # fast-forward if loaded from state
                if self._iterations_this_epoch > idx:
                    continue
                self.iterations += 1
                self._iterations_this_epoch += 1
                batch = Batch(minibatch, self.device, self.is_test)

                yield batch
            return


class TextDataloader(object):
    def __init__(self, args, datasets, batch_size,
                 device, shuffle, is_test):
        self.args = args
        self.batch_size = batch_size
        self.device = device

    def data(self):
        if self.shuffle:
            random.shuffle(self.dataset)
        xs = self.dataset
        return xs

    def preprocess(self, ex, is_test):

        src = ex['src']
        tgt = ex['tgt'][:self.args.max_tgt_len][:-1] + [2]
        src_sent_labels = ex['src_sent_labels']
        segs = ex['segs']
        if (not self.args.use_interval):
            segs = [0] * len(segs)
        clss = ex['clss']
        src_txt = ex['src_txt']
        tgt_txt = ex['tgt_txt']

        end_id = [src[-1]]
        src = src[:-1][:self.args.max_pos - 1] + end_id
        segs = segs[:self.args.max_pos]
        max_sent_id = bisect.bisect_left(clss, self.args.max_pos)
        src_sent_labels = src_sent_labels[:max_sent_id]
        clss = clss[:max_sent_id]
        # src_txt = src_txt[:max_sent_id]

        if (is_test):
            return src, tgt, segs, clss, src_sent_labels, src_txt, tgt_txt
        else:
            return src, tgt, segs, clss, src_sent_labels

    def batch_buffer(self, data, batch_size):
        minibatch, size_so_far = [], 0
        for ex in data:
            if (len(ex['src']) == 0):
                continue
            ex = self.preprocess(ex, self.is_test)
            if (ex is None):
                continue
            minibatch.append(ex)
            size_so_far = simple_batch_size_fn(ex, len(minibatch))
            if size_so_far == batch_size:
                yield minibatch
                minibatch, size_so_far = [], 0
            elif size_so_far > batch_size:
                yield minibatch[:-1]
                minibatch, size_so_far = minibatch[-1:], simple_batch_size_fn(ex, 1)
        if minibatch:
            yield minibatch

    def create_batches(self):
        """ Create batches """
        data = self.data()
        for buffer in self.batch_buffer(data, self.batch_size * 300):
            if (self.args.task == 'abs'):
                p_batch = sorted(buffer, key=lambda x: len(x[2]))
                p_batch = sorted(p_batch, key=lambda x: len(x[1]))
            else:
                p_batch = sorted(buffer, key=lambda x: len(x[2]))
                p_batch = batch(p_batch, self.batch_size)

            p_batch = batch(p_batch, self.batch_size)

            p_batch = list(p_batch)
            if (self.shuffle):
                random.shuffle(p_batch)
            for b in p_batch:
                if (len(b) == 0):
                    continue
                yield b

    def __iter__(self):
        while True:
            self.batches = self.create_batches()
            for idx, minibatch in enumerate(self.batches):
                # fast-forward if loaded from state
                if self._iterations_this_epoch > idx:
                    continue
                self.iterations += 1
                self._iterations_this_epoch += 1
                batch = Batch(minibatch, self.device, self.is_test)

                yield batch
            return
